open Ast;;

let locNum = ref 0;;

(* create a location. Increases the numbering of the locations. *)
let createLocation () = locNum := (!locNum + 1); Location(!locNum);;

(* This module handles the locations of a program *)

(* create equivalence classes *)

(* create a list of pairs on the form (l, loc) where l is a label contained in the list equiv *)
let createLabelMap loc equiv = List.map (fun l -> (l, loc)) equiv;;

(* The revers of the above: create a list of one pair (loc, equiv) where equiv is a list of labels. *)
let createLocationMap loc equiv = [ (loc, equiv) ];;

(* create a pair of a labelMap (pairs on form (l, loc)) and a locationMap (pairs on form (loc, [l])) *)
let createMapsFromLoc loc equiv = (createLabelMap loc equiv, createLocationMap loc equiv);;

(* Returns a pair of two maps given a equivalence class, i.e. this function creates a new location with maps. *)
let createMaps equiv = match equiv with
	[] -> ([],[])
	| _ -> let loc = createLocation () in (createLabelMap loc equiv, createLocationMap loc equiv);;

(* Compares two locations. Returns 0 if equal, -1 if first is less than second and 1 if first is greater than second *)
let compareLocations (Location(l1)) (Location(l2)) = compare l1 l2;;

(* Compares two locationMap's pair. *)
let compareLocMap (loc1, _) (loc2, _) = compareLocations loc1 loc2;;
	
(* Create maps from statements (the maps "this far") and a equivalence class (list of labels) returns a triple. *)
let rec createLocationsFromStm (labMap, locMap) equiv = function
											(Skip(Labels(pre, post)) | Callvoid(Labels(pre, post), _, _) | Callbool(Labels(pre, post),_,_,_) | Assignment(Labels(pre, post),_,_) | Returnvoid(Labels(pre, post)) | Returnbool(Labels(pre, post),_)) -> 
												let (labMap1, locMap1) = createMaps (pre :: equiv) in 
													(labMap @ labMap1, locMap @ locMap1, [ post ])
											| Comp(Labels(pre, post), s1, s2) -> 
												let (labMap1, locMap1, equiv1) = (createLocationsFromStm (labMap, locMap) (pre :: equiv) s1) in
												let (labMap2, locMap2, equiv2) = (createLocationsFromStm (labMap1, locMap1) equiv1 s2) in
													(labMap2, locMap2, (post :: equiv2))
											| While(Labels(pre, post),_,s) -> 
												let cm = createMapsFromLoc (createLocation ()) in
												let (labMap1, locMap1, equiv1) = createLocationsFromStm (labMap, locMap) [] s in
												let (labMap2, locMap2) = cm (pre :: (equiv @ equiv1)) in
													(labMap1 @ labMap2, locMap1 @ locMap2, [ post ] )
											| Ifelse(Labels(pre, post), e, s1, s2) -> 
												let (labMap1, locMap1) = createMaps (pre :: equiv) in 
												let (labMap2, locMap2, equiv1) = createLocationsFromStm (labMap1 @ labMap, locMap1 @ locMap) [] s1 in
												let (labMap3, locMap3, equiv2) = createLocationsFromStm (labMap2, locMap2) [] s2 in
													(labMap3, locMap3, (post :: (equiv1 @ equiv2)))
											;;
													
(* Same as above but for declarations *)
let rec createLocationsFromDecl (labMap, locMap) equiv = function
											Declaration(Labels(pre, post), _) ->
												let (labMap1, locMap1) = createMaps (pre :: equiv) in 
													(labMap @ labMap1, locMap @ locMap1, [ post ])
											| Declcomp(Labels(pre, post), d1, d2) ->
												let (labMap1, locMap1, equiv1) = createLocationsFromDecl (labMap, locMap) (pre :: equiv) d1 in
												let (labMap2, locMap2, equiv2) = createLocationsFromDecl (labMap1, locMap1) equiv1 d2 in
													(labMap2, locMap2, (post :: equiv2))
											| Empty -> (labMap, locMap, equiv)
											;;

(* Creates a new pair of a labelMap and a locMap (locationMap) given a procedure and a pair of maps. 
  The new pair of maps contains the same maps as the old one (the input) put also all maps for the procedure given. 
  labelMap = [ (label, location) ] and locMap = [ (locations, [labels]) ] *)
let createLocationsFromProc (labMap, locMap) proc = let Procedure(Labels(pre, post), _, _, _, decl, stm) = proc in
	let (labMap1, locMap1, equiv) = createLocationsFromDecl (labMap, locMap) [ pre ] decl in (* create locations for the procedure declarations *)
	let (labMap3, locMap3, equiv1) = createLocationsFromStm (labMap1, locMap1) equiv stm in (* create locations for the statements in the procedure *)
	let (labMap4, locMap4) = createMaps (post :: equiv1) in (* and for the last post-label of the statements *)
		(labMap4 @ labMap3, locMap4 @ locMap3);; (* return the pair *)
	
(* Create a pair of the maps above given a Boolean program (Bp). *)
let createLocations bp = 	
  let Bp(decl, procs) = bp in
  let rec createLocationsFromProcs (labMap, locMap) = function
	p::ps -> let (labMap1, locMap1) = createLocationsFromProc (labMap, locMap) p in createLocationsFromProcs (labMap1, locMap1)  ps
	|[] -> (labMap, locMap)
	in
  match decl with
  Empty ->
    createLocationsFromProcs ([], []) procs
  |_ ->
    let (labMap, locMap, equiv) = createLocationsFromDecl ([], []) [] decl in
    let (labMap1, locMap1) = createMaps equiv in
      createLocationsFromProcs (labMap @ labMap1, locMap @ locMap1) procs
  ;;

(* Get location from label given a label map and a label *)								
let getLocation label labMap = List.assoc label labMap;;

(* Get labels from location given a location map and a location *)								
let getLabels loc locMap = List.assoc loc locMap;;